/***************************************************************************
 *
 *   Copyright (C) 2025 by Will Gauvin
 *   Licensed under the Academic Free License version 2.1
 *
 ***************************************************************************/

#include "dsp/RescaleScaleOffsetDump.h"
#include "dsp/ASCIIObservation.h"
#include "dsp/DADAHeader.h"

#include <fstream>
#include <cstring>
#include <algorithm>

using namespace std;

dsp::RescaleScaleOffsetDump::RescaleScaleOffsetDump(const std::string& filename)
{
  FILE* ptr = std::fopen(filename.c_str(), "w");
  if (!ptr)
  {
    throw Error(FailedSys, "dsp::RescaleScaleOffsetDump::RescaleScaleOffsetDump - ctor",
        "fopen(%s)", filename.c_str());
  }
  output = ptr;
}

void dsp::RescaleScaleOffsetDump::write_header (dsp::Rescale* rescale)
{
  if (dsp::Operation::verbose)
    cerr << "dsp::RescaleScaleOffsetDump::write_header" << endl;

  const auto input = rescale->get_input();
  ASCIIObservation ascii (input);
  ascii.set_machine ("dspsr");
  // using float32 for the data.
  ascii.set_nbit(-32);
  // we have scale and offset as the dimensions
  ascii.set_ndim(2);

  // ensure we have the correct RESOLUTION, BLOCK_HEADER_BYTES and BLOCK_DATA_BYTES
  auto nchan = input->get_nchan();
  auto npol = input->get_npol();
  // 2 dims to account for the scale and the offset
  static constexpr unsigned ndim = 2;

  auto block_header_bytes = sizeof(uint64_t);
  auto block_data_bytes = nchan * npol * ndim * sizeof(float);
  auto resolution = block_header_bytes + block_data_bytes;

  DADAHeader header;
  header.resize();

  char *buffer = header.get_header();
  ascii.unload(buffer);

  if (ascii_header_set (buffer, "HDR_SIZE", "%d", header.size()) < 0)
    throw Error (InvalidState, "dsp::RescaleScaleOffsetDump::write_header",
		 "failed to set HDR_SIZE in output file header");

  if (ascii_header_set (buffer, "BLOCK_HEADER_BYTES", "%d", block_header_bytes) < 0)
    throw Error (InvalidState, "dsp::RescaleScaleOffsetDump::write_header",
		 "failed to set BLOCK_HEADER_BYTES in output file header");

  if (ascii_header_set (buffer, "BLOCK_DATA_BYTES", "%d", block_data_bytes) < 0)
    throw Error (InvalidState, "dsp::RescaleScaleOffsetDump::write_header",
		 "failed to set BLOCK_DATA_BYTES in output file header");

  if (ascii_header_set (buffer, "RESOLUTION", "%d", resolution) < 0)
    throw Error (InvalidState, "dsp::RescaleScaleOffsetDump::write_header",
		 "failed to set RESOLUTION in output file header");

  std::fwrite (buffer, sizeof(char), header.size(), output);
  header_written = true;
}


void dsp::RescaleScaleOffsetDump::handle_scale_offset_updated(dsp::Rescale* rescale)
{
  if (dsp::Operation::verbose)
    cerr << "dsp::RescaleScaleOffsetDump::handle_scale_offset_updated" << endl;

  if (!output)
  {
    throw Error(InvalidState, "dsp::RescaleScaleOffsetDump::handle_scale_offset_updated",
      "output is a nullptr. Perhaps set_output has not been called with an open file");
  }

  // if we haven't written the header -> write the given header
  if (!header_written)
    write_header(rescale);

  const unsigned nchan = rescale->get_input()->get_nchan();
  const unsigned npol = rescale->get_input()->get_npol();

  uint64_t buffer_size = sizeof(uint64_t) + 2 * nchan * npol * sizeof(float);
  if (buffer.size() < buffer_size)
  {
    buffer.resize(buffer_size);
    std::fill(buffer.begin(), buffer.end(), 0);
  }

  // want to make sure this is from the sample that matches with the output sample from when the new scales & offset were applied
  const int64_t isample = rescale->get_input()->get_input_sample();

  uint64_t idx = 0;
  memcpy(&buffer[0], &isample, sizeof(int64_t));
  idx += sizeof(int64_t);

  // order the data as [chan][pol]
  for (auto ichan = 0; ichan < nchan; ichan++)
  {
    for (auto ipol = 0; ipol < npol; ipol++)
    {
      memcpy(&buffer[idx], &rescale->get_offset(ipol)[ichan], sizeof(float));
      idx += sizeof(float);
      memcpy(&buffer[idx], &rescale->get_scale(ipol)[ichan], sizeof(float));
      idx += sizeof(float);
    }
  }

  std::fwrite(buffer.data(), sizeof(char), buffer_size, output);
  std::fflush(output);
}
